# ------------------------------------------------------------------------------
# Check exclusion restriction
# Author: Cassidy Shubatt <cshubatt@gmail.com>
# To run: bash 09_check_excl_restriction.sh
# ------------------------------------------------------------------------------

# Libraries --------------------------------------------------------------------
tictoc::tic()

library(here)
library(yaml)
library(data.table)
library(tidyverse)
library(testit) # assert()
library(glue) # glue strings
library(lfe)
library(stargazer)

temp <- here::here("code", "05_natural_experiment", "temp")
overnight_lab <- ""
u <- modules::use(here::here("lib", "util.R"))

# Load Data --------------------------------------------------------------------
message("Loading data...")
paths <- read_yaml(here::here("lib", "filepaths.yml"))
tests_in_k_hours <- readRDS(file.path(temp, "tests_in_k_hours_df.rds"))
cohort <- readRDS(glue(paths$analysis$full_cohort)) %>%
  filter(!exclude & split != "val") %>%
  u$safe_left_join(tests_in_k_hours) %>%
  mutate(hour = factor(hour)) %>%
  mutate(yhat = p__ensemble__stent_or_cabg_010_day__tested)

# Check Exclusion --------------------------------------------------------------
message("Checking exclusion restriction...")
# Exclusion concern is that maybe there are testing quotas, so shift test rate
# could influence individual test decisions
# We check by seeing if num tests in previous k hours predicts testing decision
k_list <- c(12,24,48)
excl_check_reg_list <- list()
for(k in k_list){
  prev_k_varname <- glue("ntests_in_prev_{k}")

  form <- reformulate(
    response = "test_010_day",
    termlabels = c(
      "yhat", prev_k_varname, "1 | hour + wday + week + year | 0 | ptid"
    )
  )

  fit <- felm(form, data = cohort)
  excl_check_reg_list <- c(excl_check_reg_list, list(fit))
}

# Save -------------------------------------------------------------------------
message("Saving stargazer regression results...")
reg_tbl <- stargazer(
  excl_check_reg_list, omit.stat = c("f", "ser", "ll", "aic", "adj.rsq"),
  dep.var.caption = ""
)
save_fp <- file.path(temp, "exclusion_restriction_regs.tex")
write(reg_tbl, save_fp, append = FALSE)

message("Done.")
